import mindspore.nn as nn
from mindspore.common import initializer as init
import mindspore.ops as op


class PNet(nn.Cell):
    """
    PNet of MTCNN
    """

    def __init__(self):
        super(PNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=10, kernel_size=(3, 3), has_bias=True, padding=0,
                               weight_init=init.HeNormal(mode='fan_out', nonlinearity='relu'), pad_mode='pad')  # 10,12,12
        self.prelu1 = nn.PReLU()
        # 这里计算输出形状只能向下取整  10, 6, 6
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=10, out_channels=16, kernel_size=(3, 3), has_bias=True, padding=0,
                               weight_init=init.HeNormal(mode='fan_out', nonlinearity='relu'), pad_mode='pad')
        self.prelu2 = nn.PReLU()
        self.conv3 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=(3, 3), has_bias=True, padding=0,
                               weight_init=init.HeNormal(mode='fan_out', nonlinearity='relu'), pad_mode='pad')
        self.prelu3 = nn.PReLU()
        self.conv4_1 = nn.Conv2d(in_channels=32, out_channels=2, kernel_size=(1, 1), has_bias=True, padding=0,
                                 weight_init=init.HeNormal(mode='fan_out', nonlinearity='relu'), pad_mode='pad')
        self.conv4_2 = nn.Conv2d(in_channels=32, out_channels=4, kernel_size=(1, 1), has_bias=True, padding=0,
                                 weight_init=init.HeNormal(mode='fan_out', nonlinearity='relu'), pad_mode='pad')
        self.conv4_3 = nn.Conv2d(in_channels=32, out_channels=10, kernel_size=(1, 1), has_bias=True, padding=0,
                                 weight_init=init.HeNormal(mode='fan_out', nonlinearity='relu'), pad_mode='pad')

        self.squeeze = op.Squeeze(axis=2)

    def construct(self, x):
        x = self.prelu1(self.conv1(x))  # (1, 10, 10, 10)
        x = self.pool1(x)  # (1, 10, 6, 6)
        x = self.prelu2(self.conv2(x))  # (1, 16, 3, 3)
        x = self.prelu3(self.conv3(x))  # (1, 32, 1, 1)
        # 分类是否人脸的卷积输出层
        class_out = self.conv4_1(x)  # (1, 2, 1, 1)
        class_out = self.squeeze(class_out)
        class_out = self.squeeze(class_out)
        # 人脸box的回归卷积输出层
        bbox_out = self.conv4_2(x)  # (1, 4, 1, 1)
        bbox_out = self.squeeze(bbox_out)
        bbox_out = self.squeeze(bbox_out)
        # 5个关键点的回归卷积输出层
        landmark_out = self.conv4_3(x)  # (1, 10, 1, 1)
        landmark_out = self.squeeze(landmark_out)
        landmark_out = self.squeeze(landmark_out)
        return class_out, bbox_out, landmark_out
